import torch
from torch import nn

"""
CIFAR10数据集的网络模型，参考图片 CIFAR10分类模型.png
"""
class AdamModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.model = nn.Sequential(
            nn.Conv2d(in_channels=3, out_channels=32, kernel_size=5, stride=1, padding=2), # 卷积
            nn.MaxPool2d(2),# 最大池化，减小宽高尺寸， 最大池化，相当于是将图片压缩了， 在保留图片特征的情况下减少了像素
            nn.Conv2d(in_channels=32, out_channels=32, kernel_size=5, stride=1, padding=2),
            nn.MaxPool2d(2),
            nn.Conv2d(in_channels=32, out_channels=64, kernel_size=5, stride=1, padding=2),
            nn.MaxPool2d(2),
            nn.Flatten(),
            nn.Linear(in_features=64 * 4 * 4, out_features=64),
            nn.Linear(in_features=64, out_features=10)
        )

    def forward(self, input):
        return self.model.forward(input)

if __name__ == '__main__':
    model = AdamModel()
    input = torch.ones((64, 3, 32, 32)) # 64个数据， 3通道，宽32, 高32
    output = model(input)
    print(output.shape) # 应该是长度为10
